import sqlite3
from typing import Dict, List, Optional
import re


class TCMDBHandler:
    """
    中医数据库处理器类（适用于SQLite）
    功能：管理数据库连接，封装中药材查询方法，为TCMLLM提供数据支持
    """

    def __init__(self, db_path: str):
        """
        初始化数据库连接（SQLite通过文件路径连接）
        :param db_path: SQLite数据库文件路径（如：'tcm.db'）
        """
        self.db_path = db_path
        self.connection = sqlite3.connect(db_path)
        # 启用外键约束（SQLite默认关闭）
        self.connection.execute("PRAGMA foreign_keys = ON;")
        # 设置查询结果为字典格式
        self.connection.row_factory = sqlite3.Row

    def __enter__(self):
        """支持上下文管理器（with语句），自动管理连接"""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """退出时关闭数据库连接"""
        self.connection.close()

    def _execute_query(self, sql: str, params: tuple = ()) -> List[Dict]:
        """
        执行通用查询方法（内部使用）
        :param sql: SQL查询语句
        :param params: 查询参数（防止SQL注入）
        :return: 查询结果列表（字典形式，键为字段名）
        """
        cursor = self.connection.cursor()
        cursor.execute(sql, params)
        # 将sqlite3.Row转换为字典
        results = [dict(row) for row in cursor.fetchall()]
        cursor.close()
        return results

    def get_herb_by_name(self, herb_name: str) -> Optional[Dict]:
        """
        按药材名称精确查询
        :param herb_name: 药材名称（如："黄芪"）
        :return: 药材信息字典（未找到时返回None）
        """
        sql = "SELECT * FROM tcm_herb WHERE herb_name = ?;"
        results = self._execute_query(sql, (herb_name,))
        return results[0] if results else None

    def get_herb_by_alias(self, alias: str) -> List[Dict]:
        """
        按别名模糊查询（支持部分匹配）
        :param alias: 别名关键词（如："棉芪"会匹配包含该词的所有别名）
        :return: 匹配的药材信息列表（可能多个）
        """
        sql = "SELECT * FROM tcm_herb WHERE alias LIKE ?;"
        return self._execute_query(sql, (f"%{alias}%",))

    def search_herb_by_property(self, property_type: str, keyword: str) -> List[Dict]:
        """
        按药性属性（功效、主治、性味、归经）模糊搜索
        :param property_type: 属性字段名（可选值：'efficacy', 'indication', 'property', 'meridian'）
        :param keyword: 搜索关键词（如："补气", "健脾"）
        :return: 匹配的药材信息列表
        """
        valid_properties = {"efficacy", "indication", "property", "meridian"}
        if property_type not in valid_properties:
            raise ValueError(
                "无效的属性类型！有效值：efficacy, indication, property, meridian"
            )
        sql = f"SELECT * FROM tcm_herb WHERE {property_type} LIKE ?;"
        return self._execute_query(sql, (f"%{keyword}%",))

    def format_herb_info(self, herb_data: Dict) -> str:
        """
        将数据库查询结果格式化为LLM可直接使用的回答文本
        :param herb_data: 药材信息字典（来自数据库查询结果）
        :return: 结构化的中文回答字符串
        """
        return (
            f"【药材名称】{herb_data['herb_name']}\n"
            f"【别名】{herb_data['alias'] if herb_data['alias'] else '无'}\n"
            f"【性味】{herb_data['property'] if herb_data['property'] else '无'}\n"
            f"【归经】{herb_data['meridian'] if herb_data['meridian'] else '无'}\n"
            f"【功效】{herb_data['efficacy'] if herb_data['efficacy'] else '无相关记载'}\n"
            f"【主治】{herb_data['indication'] if herb_data['indication'] else '无相关记载'}\n"
            f"【用法用量】{herb_data['usage'] if herb_data['usage'] else '无明确记载'}\n"
            f"【禁忌】{herb_data['contraindication'] if herb_data['contraindication'] else '无特殊禁忌'}"
        )


class TCMLLM:
    def __init__(self, db_path: str):
        self.db = TCMDBHandler(db_path)
        # 新增：别名联想规则（解决用户用别名提问的问题）
        self.alias_map = self._load_alias_map()  # 从数据库加载别名映射表
        # 定义用户问题→数据库字段的映射规则（正则表达式匹配）
        self.intent_rules = [
            (r"^(.+?)的?功效$", "efficacy"),
            (r"^(.+?)的?主治$", "indication"),
            (r"^(.+?)的?性味$", "property"),
            (r"^(.+?)的?归经$", "meridian"),
            (r"^(.+?)的?怎么用$", "usage"),
            (r"^(.+?)的?禁忌$", "contraindication"),
            (r"^查找(.*)$", "full_info"),
        ]

    def _load_alias_map(self) -> Dict[str, str]:
        """从数据库预加载别名→正式名称的映射（提升查询成功率）"""
        sql = "SELECT herb_name, alias FROM tcm_herb WHERE alias IS NOT NULL;"
        alias_map = {}
        for row in self.db.connection.execute(sql).fetchall():
            for alias in row["alias"].split(";"):
                alias_map[alias.strip()] = row["herb_name"]
        return alias_map

    def _parse_user_query(self, query: str):
        """增强：将别名转换为正式药材名称，并解析查询意图"""
        query = self.alias_map.get(query, query)  # 替换别名→正式名
        for pattern, field in self.intent_rules:
            match = re.match(pattern, query.strip())
            if match:
                herb_name = match.group(1).strip()
                return {"herb_name": herb_name, "target_field": field}
        # 未匹配到规则时，默认查询完整信息
        return {"herb_name": query, "target_field": "full_info"}

    def _generate_answer(self, herb_data: Dict, target_field: str) -> str:
        """根据查询意图生成回答"""
        if target_field == "full_info":
            return self.db.format_herb_info(herb_data)
        # 单字段查询（如只返回功效/主治）
        value = herb_data.get(target_field, "无相关记载")
        return f"{herb_data['herb_name']}的{target_field.replace('_', ' ')}：{value}"

    def process_query(self, user_query: str) -> str:
        """LLM核心处理流程：解析→查询→生成回答"""
        # 1. 解析用户意图
        parsed = self._parse_user_query(user_query)
        herb_name = parsed["herb_name"]
        target_field = parsed["target_field"]

        # 2. 数据库查询
        herb_data = self.db.get_herb_by_name(herb_name)
        if not herb_data:
            return f"未找到药材【{herb_name}】的信息，请检查名称是否正确。"

        # 3. 生成回答
        return self._generate_answer(herb_data, target_field)


if __name__ == "__main__":
    llm = TCMLLM("tcm.db")
    user_question = "黄芪的功效"
    response = llm.process_query(user_question)
    print("LLM回答：\n", response)
